package com.ruoyi.common.utils.taskscheduler;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.*;

/**
 * 定时任务调度器，支持三种类型的任务：
 * 1. 一次性延时任务
 * 2. 固定速率周期性任务
 * 3. 固定延迟周期性任务
 */
public class TaskScheduler {

    private final ScheduledThreadPoolExecutor scheduler;

    private final Map<String, ScheduledFuture<?>> taskMap;

    private final Object lock;

    public TaskScheduler(int corePoolSize) {
        this.scheduler = new ScheduledThreadPoolExecutor(
                corePoolSize,
                Executors.defaultThreadFactory(),
                new ThreadPoolExecutor.DiscardPolicy() // 任务太多是采用丢弃策略
        );

        this.taskMap = new HashMap<>();
        this.lock = new Object();
    }

    /**
     * 添加一次性延时任务
     * @param taskId 任务ID
     * @param task 具体要执行的任务
     * @param delay 延迟时间
     * @param unit 时间单位
     * @return 是否添加成功
     */
    public boolean scheduleOneTimeTask(String taskId, Runnable task, long delay, TimeUnit unit) {
        synchronized (lock) {
            if (taskMap.containsKey(taskId)) {
                return false;
            }
            ScheduledFuture<?> future = scheduler.schedule(task, delay, unit);
            taskMap.put(taskId, future);

            // 添加任务完成后从map中移除
            scheduler.schedule(() -> {
                synchronized (lock) {
                    taskMap.remove(taskId);
                }
            },delay, unit);
            return true;
        }
    }

    /**
     * 添加固定速率周期性任务
     * @param taskId 任务唯一标识
     * @param task 要执行的任务
     * @param initialDelay 初始延迟时间
     * @param period 执行周期
     * @param unit 时间单位
     * @return 是否添加成功
     */
    public boolean scheduleFixedRateTask(String taskId, Runnable task, long initialDelay,
                                         long period, TimeUnit unit) {
        synchronized (lock) {
            if (taskMap.containsKey(taskId)) {
                return false; // 任务ID已存在
            }

            ScheduledFuture<?> future = scheduler.scheduleAtFixedRate(task, initialDelay, period, unit);
            taskMap.put(taskId, future);
            return true;
        }
    }

    /**
     * 添加固定延迟周期性任务
     * @param taskId 任务唯一标识
     * @param task 要执行的任务
     * @param initialDelay 初始延迟时间
     * @param delay 任务执行完成后的延迟时间
     * @param unit 时间单位
     * @return 是否添加成功
     */
    public boolean scheduleFixedDelayTask(String taskId, Runnable task, long initialDelay,
                                          long delay, TimeUnit unit) {
        synchronized (lock) {
            if (taskMap.containsKey(taskId)) {
                return false; // 任务ID已存在
            }

            ScheduledFuture<?> future = scheduler.scheduleWithFixedDelay(task, initialDelay, delay, unit);
            taskMap.put(taskId, future);
            return true;
        }
    }

    /**
     * 添加带返回值的一次性延时任务
     * @param taskId 任务唯一标识
     * @param task 要执行的任务
     * @param delay 延迟时间
     * @param unit 时间单位
     * @return 是否添加成功
     */
    public ScheduledFuture<String> scheduleOneTimeCallable(String taskId, Callable<String> task, long delay, TimeUnit unit) {
        synchronized (lock) {
            if (taskMap.containsKey(taskId)) {
                return null;
            }

            ScheduledFuture<String> future = scheduler.schedule(task, delay, unit);
            taskMap.put(taskId, future);

            // 任务完成后从map中移除
            scheduler.schedule(() -> {
                synchronized (lock) {
                    taskMap.remove(taskId);
                }
            }, delay, unit);

            return future;
        }
    }

    /**
     * 取消任务
     * @param taskId 任务唯一标识
     * @param mayInterruptIfRunning 是否允许中断正在执行的任务
     * @return 是否取消成功
     */
    public boolean cancelTask(String taskId, boolean mayInterruptIfRunning) {
        synchronized (lock) {
            ScheduledFuture<?> future = taskMap.get(taskId);
            if (future == null) {
                return false;
            }

            boolean result = future.cancel(mayInterruptIfRunning);
            if (result) {
                taskMap.remove(taskId);
            }
            return result;
        }
    }

    /**
     * 查询任务是否已完成
     * @param taskId 任务唯一标识
     * @return 任务是否已完成，如果任务不存在返回null
     */
    public Boolean isTaskDone(String taskId) {
        synchronized (lock) {
            ScheduledFuture<?> future = taskMap.get(taskId);
            if (future == null) {
                return null;
            }
            return future.isDone();
        }
    }

    /**
     * 查询任务是否已取消
     * @param taskId 任务唯一标识
     * @return 任务是否已取消，如果任务不存在返回null
     */
    public Boolean isTaskCancelled(String taskId) {
        synchronized (lock) {
            ScheduledFuture<?> future = taskMap.get(taskId);
            if (future == null) {
                return null;
            }
            return future.isCancelled();
        }
    }

    /**
     * 获取当前活跃的任务数量
     * @return 活跃任务数量
     */
    public int getActiveTaskCount() {
        synchronized (lock) {
            return taskMap.size();
        }
    }

    /**
     * 关闭调度器
     * @param awaitTermination 是否等待所有任务完成
     * @param timeout 等待超时时间
     * @param unit 时间单位
     * @return 是否成功关闭
     * @throws InterruptedException 等待过程中被中断
     */
    public boolean shutdown(boolean awaitTermination, long timeout, TimeUnit unit) throws InterruptedException {
        scheduler.shutdown();
        if (awaitTermination) {
            return scheduler.awaitTermination(timeout, unit);
        }
        return true;
    }

    /**
     * 立即关闭调度器，尝试中断所有正在执行的任务
     */
    public void shutdownNow() {
        scheduler.shutdownNow();
        synchronized (lock) {
            taskMap.clear();
        }
    }

}
